# ******************************************************************************
#  Copyright (c) 2021-2022. Kneron Inc. All rights reserved.                   *
# ******************************************************************************

# Legacy Module
from .Legacy.V1.KPStructure import \
    GenericRawImageHeaderBuffer, \
    GenericRawResultHeaderBuffer, \
    GenericRawBypassPreProcImageHeaderBuffer, \
    GenericRawBypassPreProcResultHeaderBuffer

# Core Module
from .KPStructure import \
    DdrManageAttributesBuffer, \
    DeviceDescriptorListBuffer, \
    ModelNefDescriptorBuffer, \
    SystemInfoBuffer, \
    InfConfigurationBuffer, \
    GenericImageInferenceDescriptorBuffer, \
    GenericImageInferenceResultHeaderBuffer, \
    GenericDataInferenceImageHeaderBuffer, \
    GenericDataInferenceResultHeaderBuffer, \
    InfFixedNodeOutputBuffer, \
    InfFloatNodeOutputBuffer, \
    ProfileDataBuffer, \
    PerformanceMonitorDataBuffer
from .KPBaseClass.WrapperBase import WrapperBase
from .KPLibLoader import KPLibLoader
import ctypes


class KPWrapper(WrapperBase):
    def __init__(self):
        super(KPWrapper, self).__init__(lib_loader=KPLibLoader())

    def _init_sdk_functions(self):
        self.__init_kp_scan_devices()
        self.__init_kp_connect_devices()
        self.__init_kp_connect_devices_without_check()
        self.__init_kp_disconnect_devices()
        self.__init_kp_set_timeout()
        self.__init_kp_reset_device()
        self.__init_kp_load_firmware()
        self.__init_kp_load_firmware_from_file()
        self.__init_kp_load_model()
        self.__init_kp_load_model_from_file()
        self.__init_kp_load_encrypted_models()
        self.__init_kp_load_encrypted_models_from_file()
        self.__init_kp_enable_firmware_log()
        self.__init_kp_disable_firmware_log()
        self.__init_kp_get_system_info()
        self.__init_kp_get_model_info()
        self.__init_kp_release_model_nef_descriptor()
        self.__init_kp_update_kdp2_firmware()
        self.__init_kp_update_kdp2_firmware_from_files()
        self.__init_kp_load_model_from_flash()
        self.__init_kp_store_ddr_manage_attr()
        self.__init_kp_install_driver_for_windows()
        self.__init_kp_get_version()
        self.__init_c_free()

        self.__init_kp_inference_configure()
        self.__init_kp_generic_raw_inference_send()
        self.__init_kp_generic_raw_inference_receive()
        self.__init_kp_generic_raw_inference_bypass_pre_proc_send()
        self.__init_kp_generic_raw_inference_bypass_pre_proc_receive()
        self.__init_kp_generic_image_inference_send()
        self.__init_kp_generic_image_inference_receive()
        self.__init_kp_generic_data_inference_send()
        self.__init_kp_generic_data_inference_receive()
        self.__init_kp_generic_inference_retrieve_fixed_node()
        self.__init_kp_generic_inference_retrieve_float_node()
        self.__init_kp_profile_set_enable()
        self.__init_kp_profile_get_statistics()
        self.__init_kp_performance_monitor_set_enable()
        self.__init_kp_performance_monitor_get_statistics()

        self.__init_kp_set_ckey()
        self.__init_kp_set_secure_boot_key()
        self.__init_kp_set_gpio()

    def __init_kp_scan_devices(self):
        c_function = self._lib.kp_scan_devices
        c_function.argtypes = None
        c_function.restype = ctypes.POINTER(DeviceDescriptorListBuffer)

    def __init_kp_connect_devices(self):
        c_function = self._lib.kp_connect_devices
        c_function.argtypes = [ctypes.c_int,
                               ctypes.POINTER(ctypes.c_int),
                               ctypes.POINTER(ctypes.c_int)]
        c_function.restype = ctypes.c_void_p

    def __init_kp_connect_devices_without_check(self):
        c_function = self._lib.kp_connect_devices_without_check
        c_function.argtypes = [ctypes.c_int,
                               ctypes.POINTER(ctypes.c_int),
                               ctypes.POINTER(ctypes.c_int)]
        c_function.restype = ctypes.c_void_p

    def __init_kp_disconnect_devices(self):
        c_function = self._lib.kp_disconnect_devices
        c_function.argtypes = [ctypes.c_void_p]
        c_function.restype = ctypes.c_int

    def __init_kp_set_timeout(self):
        c_function = self._lib.kp_set_timeout
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_int]
        c_function.restype = None

    def __init_kp_reset_device(self):
        c_function = self._lib.kp_reset_device
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_int]
        c_function.restype = ctypes.c_int

    def __init_kp_load_firmware(self):
        c_function = self._lib.kp_load_firmware
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_void_p,
                               ctypes.c_int,
                               ctypes.c_void_p,
                               ctypes.c_int]
        c_function.restype = ctypes.c_int

    def __init_kp_load_firmware_from_file(self):
        c_function = self._lib.kp_load_firmware_from_file
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_char_p,
                               ctypes.c_char_p]
        c_function.restype = ctypes.c_int

    def __init_kp_load_model(self):
        c_function = self._lib.kp_load_model
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_void_p,
                               ctypes.c_int,
                               ctypes.POINTER(ModelNefDescriptorBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_load_model_from_file(self):
        c_function = self._lib.kp_load_model_from_file
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_char_p,
                               ctypes.POINTER(ModelNefDescriptorBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_load_encrypted_models(self):
        c_function = self._lib.kp_load_encrypted_models
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(ctypes.c_void_p),
                               ctypes.c_int,
                               ctypes.c_int,
                               ctypes.POINTER(ModelNefDescriptorBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_load_encrypted_models_from_file(self):
        c_function = self._lib.kp_load_encrypted_models_from_file
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(ctypes.c_char_p),
                               ctypes.c_int,
                               ctypes.POINTER(ModelNefDescriptorBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_enable_firmware_log(self):
        c_function = self._lib.kp_enable_firmware_log
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_int,
                               ctypes.c_char_p]
        c_function.restype = ctypes.c_int

    def __init_kp_disable_firmware_log(self):
        c_function = self._lib.kp_disable_firmware_log
        c_function.argtypes = [ctypes.c_void_p]
        c_function.restype = ctypes.c_int

    def __init_kp_get_system_info(self):
        c_function = self._lib.kp_get_system_info
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_int,
                               ctypes.POINTER(SystemInfoBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_get_model_info(self):
        c_function = self._lib.kp_get_model_info
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_int,
                               ctypes.POINTER(ModelNefDescriptorBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_release_model_nef_descriptor(self):
        c_function = self._lib.kp_release_model_nef_descriptor
        c_function.argtypes = [ctypes.POINTER(ModelNefDescriptorBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_update_kdp2_firmware(self):
        c_function = self._lib.kp_update_kdp2_firmware
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_void_p,
                               ctypes.c_int,
                               ctypes.c_void_p,
                               ctypes.c_int,
                               ctypes.c_bool]
        c_function.restype = ctypes.c_int

    def __init_kp_update_kdp2_firmware_from_files(self):
        c_function = self._lib.kp_update_kdp2_firmware_from_files
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_char_p,
                               ctypes.c_char_p,
                               ctypes.c_bool]
        c_function.restype = ctypes.c_int

    def __init_kp_load_model_from_flash(self):
        c_function = self._lib.kp_load_model_from_flash
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(ModelNefDescriptorBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_store_ddr_manage_attr(self):
        c_function = self._lib.kp_store_ddr_manage_attr
        c_function.argtypes = [ctypes.c_void_p,
                               DdrManageAttributesBuffer]
        c_function.restype = ctypes.c_int

    def __init_kp_install_driver_for_windows(self):
        c_function = self._lib.kp_install_driver_for_windows
        c_function.argtypes = [ctypes.c_int]
        c_function.restype = ctypes.c_int

    def __init_kp_get_version(self):
        c_function = self._lib.kp_get_version
        c_function.argtypes = None
        c_function.restype = ctypes.c_char_p

    def __init_c_free(self):
        c_function = self._lib.py_c_free
        c_function.argtypes = [ctypes.c_void_p]
        c_function.restype = None

    def __init_kp_inference_configure(self):
        c_function = self._lib.kp_inference_configure
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(InfConfigurationBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_generic_raw_inference_send(self):
        c_function = self._lib.kp_generic_raw_inference_send
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(GenericRawImageHeaderBuffer),
                               ctypes.POINTER(ctypes.c_uint8)]
        c_function.restype = ctypes.c_int

    def __init_kp_generic_raw_inference_receive(self):
        c_function = self._lib.kp_generic_raw_inference_receive
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(GenericRawResultHeaderBuffer),
                               ctypes.POINTER(ctypes.c_uint8),
                               ctypes.c_uint32]
        c_function.restype = ctypes.c_int

    def __init_kp_generic_raw_inference_bypass_pre_proc_send(self):
        c_function = self._lib.kp_generic_raw_inference_bypass_pre_proc_send
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(GenericRawBypassPreProcImageHeaderBuffer),
                               ctypes.POINTER(ctypes.c_uint8)]
        c_function.restype = ctypes.c_int

    def __init_kp_generic_raw_inference_bypass_pre_proc_receive(self):
        c_function = self._lib.kp_generic_raw_inference_bypass_pre_proc_receive
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(GenericRawBypassPreProcResultHeaderBuffer),
                               ctypes.POINTER(ctypes.c_uint8),
                               ctypes.c_uint32]
        c_function.restype = ctypes.c_int

    def __init_kp_generic_image_inference_send(self):
        c_function = self._lib.kp_generic_image_inference_send
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(GenericImageInferenceDescriptorBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_generic_image_inference_receive(self):
        c_function = self._lib.kp_generic_image_inference_receive
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(GenericImageInferenceResultHeaderBuffer),
                               ctypes.POINTER(ctypes.c_uint8),
                               ctypes.c_uint32]
        c_function.restype = ctypes.c_int

    def __init_kp_generic_data_inference_send(self):
        c_function = self._lib.kp_generic_data_inference_send
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(GenericDataInferenceImageHeaderBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_generic_data_inference_receive(self):
        c_function = self._lib.kp_generic_data_inference_receive
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(GenericDataInferenceResultHeaderBuffer),
                               ctypes.POINTER(ctypes.c_uint8),
                               ctypes.c_uint32]
        c_function.restype = ctypes.c_int

    def __init_kp_generic_inference_retrieve_fixed_node(self):
        c_function = self._lib.kp_generic_inference_retrieve_fixed_node
        c_function.argtypes = [ctypes.c_uint32,
                               ctypes.POINTER(ctypes.c_uint8),
                               ctypes.c_int]
        c_function.restype = ctypes.POINTER(InfFixedNodeOutputBuffer)

    def __init_kp_generic_inference_retrieve_float_node(self):
        c_function = self._lib.kp_generic_inference_retrieve_float_node
        c_function.argtypes = [ctypes.c_uint32,
                               ctypes.POINTER(ctypes.c_uint8),
                               ctypes.c_int]
        c_function.restype = ctypes.POINTER(InfFloatNodeOutputBuffer)

    def __init_kp_profile_set_enable(self):
        c_function = self._lib.kp_profile_set_enable
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_bool]
        c_function.restype = ctypes.c_int

    def __init_kp_profile_get_statistics(self):
        c_function = self._lib.kp_profile_get_statistics
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(ProfileDataBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_performance_monitor_set_enable(self):
        c_function = self._lib.kp_performance_monitor_set_enable
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_bool]
        c_function.restype = ctypes.c_int

    def __init_kp_performance_monitor_get_statistics(self):
        c_function = self._lib.kp_performance_monitor_get_statistics
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.POINTER(PerformanceMonitorDataBuffer)]
        c_function.restype = ctypes.c_int

    def __init_kp_set_ckey(self):
        c_function = self._lib.kp_set_ckey
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_uint32]
        c_function.restype = ctypes.c_int

    def __init_kp_set_secure_boot_key(self):
        c_function = self._lib.kp_set_secure_boot_key
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_uint32,
                               ctypes.c_uint32]
        c_function.restype = ctypes.c_int

    def __init_kp_set_gpio(self):
        c_function = self._lib.kp_set_gpio
        c_function.argtypes = [ctypes.c_void_p,
                               ctypes.c_uint32,
                               ctypes.c_uint32]
        c_function.restype = ctypes.c_int
